source(here::here("code/load.R"))

#load functions in .R
source(here("code/metastudiesfunctions.R"))

#load data
dat <- read_csv(here("data/estimates.csv"))
survey_dat <- readRDS(here("data/survey_clean.rds"))

#power function
get_power <- function(std.error, truth, alpha = .05, onetail = FALSE) {
  true_z <- abs(truth) / std.error
  if (isTRUE(onetail)) {
    critical_z <- qnorm(1 - alpha)
    # probability of rejecting the null to the right
    p.hi <- 1 - pnorm(critical_z - true_z)
    power <- p.hi
  } else {
    critical_z <- qnorm(1 - alpha / 2)
    # probability of rejecting the null to the right
    p.hi <- 1 - pnorm(critical_z - true_z)
    # probability of rejecting the null to the left
    p.lo <- pnorm(-critical_z - true_z)
    power <- p.lo + p.hi
  }
  return(power)
}

##AK method, calling scripts from https://github.com/maxkasy/MetaStudiesApp/blob/master/metastudiesfunctions.r
# asymmetrical, normal, cutoffs at -1.96, 0, 1.96.
get_ak_inner <- function(b, se, cutoffs, symmetric, model = "normal") {
  estimates <- metastudies_estimation(X = b,
                                      sigma = se,
                                      cutoffs = cutoffs,
                                      symmetric = symmetric,
                                      model = model)
  out <- estimatestable(estimates$Psihat,
                        estimates$SE,
                        cutoffs = cutoffs,
                        symmetric = symmetric,
                        model = model)
  return(out[1])
}
# get_ak <- function(estimate, ...) {
#     void <- capture.output(out <- get_ak_inner(b = estimate, ...))
#     return(out)
# }

# weird but inconsequential error is generated on linux but not windows or mac, so we try-catch.
get_ak <- function(estimate, ...) {
    void <- capture.output(out <- try(get_ak_inner(b = estimate, ...), silent = TRUE))
    if (inherits(out, "try-error")) {
        out <- rep(NA, length(estimate))
    }
    return(out)
}

#add AK data
dat_ak <- dat |>
  #very many TriWen MAs or small n MAs lead to singular matrices, 
  filter(!grepl("TriWen", question_id),
         n > 9) |>
  group_by(question_id) |>
  mutate(
    ak_asym = get_ak(estimate, std.error, cutoffs = c(-1.96,1.96), symmetric = F),
    power_ak_asym = get_power(std.error, ak_asym)
    ) |>
  ungroup()


dat_ak |>
  #pivot data so power column is a variable, remove power_ prefix
  pivot_longer(cols = c("power_uwls", "power_ak_asym", "power_petpeese", "power_waapuwls"),
               names_to = "power", values_to = "power_value",
               names_prefix = "power_") |>
  ggplot(aes(x = power_value)) +
  geom_histogram() +
  facet_wrap(vars(power))

#power is estimated with much more variability for AK
dat_ak |>
  mutate(gap = power_ak_asym - power_uwls) |>
  group_by(question_id) |>
  summarize(mean_gap = mean(gap, na.rm = TRUE),
            n = n()) |>
  arrange(desc(mean_gap)) |>
  ggplot(aes(x = n, y = mean_gap)) +
  geom_point() +
  labs(y = "AK power - UWLS power",
       x = "Number of estimates") +
  scale_y_continuous(limits = c(-1, 1))

dat_ak |>
  mutate(gap = power_petpeese - power_uwls) |>
  group_by(question_id) |>
  summarize(mean_gap = mean(gap, na.rm = TRUE),
            n = n()) |>
  arrange(desc(mean_gap)) |>
  ggplot(aes(x = n, y = mean_gap)) +
  geom_point() +
  labs(y = "PET-PEESE power - UWLS power",
       x = "Number of estimates") +
  scale_y_continuous(limits = c(-1, 1))

dat_ak |>
  mutate(gap = power_waapuwls - power_uwls) |>
  group_by(question_id) |>
  summarize(mean_gap = mean(gap, na.rm = TRUE),
            n = n()) |>
  arrange(desc(mean_gap)) |>
  ggplot(aes(x = n, y = mean_gap)) +
  geom_point() +
  labs(y = "Power from WAAP - Power from UWLS",
       x = "Number of estimates") +
  scale_y_continuous(limits = c(-1, 1))


#main plot, with ak results
main_plot <- function(fn, metrics, omit = NULL) {
  fn <- here::here(sprintf("fig/power_plot%s.png", fn))
  power <- seq(0.06, .99, .01)
  get_share <- function(x) sapply(power, function(p) mean(x > p, na.rm = TRUE))
  
  if (!is.null(omit)) {
    dat_plot <- dat_ak |> filter(meta_id != {{omit}})
  } else {
    dat_plot <- dat_ak
  }
  
  dat_plot <- lapply(dat_plot[, metrics], get_share) |>
    data.frame() |>
    setNames(names(metrics)) |>
    mutate(power = power) |>
    pivot_longer(cols = seq_along(metrics)) |>
    # use factor to sort colors in the legend
    arrange(power, value) |>
    mutate(name = factor(name, name[length(metrics):1]),
           value = value * 100,
           power = power * 100)
  
  if (!is.null(omit)) {
    dat_plot$meta_id <- omit
    return(dat_plot)
  }
  
  p <- ggplot() +
    #strip plots
    geom_boxplot(survey_dat,
                 mapping = aes(x = 50, y = power_50),
                 width = 5,
                 outlier.shape = NA,
                 alpha = 0,
                 coef = 0) +
    geom_jitter(survey_dat,
                mapping = aes(x = 50, y = power_50),
                width = 1,
                alpha = 0.5) +
    geom_boxplot(survey_dat,
                 mapping = aes(x = 80, y = power_80),
                 width = 5,
                 outlier.shape = NA,
                 alpha = 0,
                 coef = 0) +
    geom_jitter(survey_dat,
                mapping = aes(x = 80, y = power_80),
                width = 1,
                alpha = 0.5) +
    geom_line(dat_plot, mapping = aes(x = power, y = value, color = name)) +
    scale_x_continuous(breaks = seq(0, 100, 20)) +
    scale_y_continuous(breaks = seq(0, 100, 20), limits = c(0, 100)) +
    theme(legend.title = element_blank()) +
    scale_color_manual(values = discrete_palette) +
    labs(x = "Power level", y = "Share of all estimates")
  ggsave(p, filename = fn, width = 6, height = 4)
}

main_plot(
  fn <- "_ak",
  metrics = c(
    "AK" = "power_ak_asym",
    "PET-PEESE" = "power_petpeese",
    "UWLS" = "power_uwls",
    "WAAP-UWLS" = "power_waapuwls"))

